import json
from langchain_community.vectorstores import Chroma
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter

from Embed_model import BGE_Embed

# 加载 wechat.txt 文件
with open('./docs/wechat.txt', 'r', encoding='utf-8') as file:
    content = file.read()

# 提取文本内容
texts = []
# 使用三引号分割内容
entries = content.split('"""')
for entry in entries:
    lines = entry.strip().split('\n')
    if len(lines) >= 2:
        question = lines[0].strip()
        answer = '\n'.join(lines[1:]).strip()
        full_text = f"{question}\n{answer}"
        texts.append(full_text)

# print(texts)

# 创建文本分割器
text_splitter = RecursiveCharacterTextSplitter(
    separators=[
        # 以"""分割文本
        "\"\"\"",
    ],
    chunk_size=1000,
    chunk_overlap=100
)

# 分割文档，获取文本块
chunks = []
for text in texts:
    chunks.extend(text_splitter.split_text(text))


# 将字符串列表转换为 Document 对象列表
documents = [Document(page_content=chunk) for chunk in chunks]

# 加载词嵌入模型
embedding = BGE_Embed()

# 创建向量仓库
vectorstore = Chroma.from_documents(
    documents=documents,
    embedding=embedding,
    persist_directory='./chroma_db'
)
